# from Banana.dev as inspiration, but PyClowder already does this.
from banana_dev import Client

# Create a reference to your model on Banana
my_model = Client(
    api_key="YOUR_API_KEY", # Found in dashboard
    model_key="YOUR_MODEL_KEY", # Found in model view in dashboard
    url="http://localhost:8000",
)

# Specify the model's input JSON
inputs = {
    "prompt": "Software developers start with a Hello, [MASK]! script",
}

# Call your model's inference endpoint
result, meta = my_model.call("/", inputs)
print(result)

# result:
# {
#     "outputs": {
#         "score": 0.8281897306442261,
#         "sequence": "software developers start with a hello, world! script.",
#         "token": 2088,
#         "token_str": "world"
#     }
# }

# https://github.com/replicate/cog
from cog import BasePredictor, Input, Path
import torch

class Predictor(BasePredictor):
    def setup(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.model = torch.load("./weights.pth")

    # The arguments and types the model takes as input
    def predict(self,
          image: Path = Input(description="Grayscale input image")
    ) -> Path:
        """Run a single prediction on the model"""
        processed_image = preprocess(image)
        output = self.model(processed_image)
        return postprocess(output)